from tensorflow.keras import Model, Input

from modelDesign import get_custom_objects, Encoder, Decoder


feedback_bits = 375
B = 3

input_shape = (24, 16, 2)
Encoder_input = Input(shape=input_shape, name="encoder_input")
Encoder_output = Encoder(Encoder_input, 128, 2, rezero=True, feedback_bits=feedback_bits, B=B, return_rx=False)
# Encoder_output = Encoder(Encoder_input, 32, 2, rezero=True, feedback_bits=feedback_bits, B=B, return_rx=False)
encoder = Model(inputs=Encoder_input, outputs=Encoder_output, name='encoder')

encoder.load_weights("../user_data/encoder_m.h5")
encoder.save("../user_data/encoder.h5")

Decoder_input = Input(shape=(feedback_bits,), name='decoder_input')
Decoder_output = Decoder(Decoder_input, 128, 27, rezero=True, feedback_bits=feedback_bits, B=B, return_rx=False)
# Decoder_output = Decoder(Decoder_input, 32, 2, rezero=True, feedback_bits=feedback_bits, B=B, return_rx=False)
decoder = Model(inputs=Decoder_input, outputs=Decoder_output, name="decoder")

decoder.load_weights("../user_data/decoder_m.h5")
decoder.save("../user_data/decoder.h5")
